import json
from elasticsearch import Elasticsearch
from langchain_openai import ChatOpenAI
from fastapi import Request
from app.const.es_output_scheme import (
    PROMQT_INFOMATION,
    OUTPUT_STRING,
    OUTPUT_STRING_1,
)
from app.config import ElasticSearchConfig
from app.exceptions.ESException import (
    es_connect_url_required_exception,
    es_index_name_required_exception,
    es_client_err_exception,
)


class ElasticiSearchPromqtManager:

    def __init__(self, request: Request):
        self.es_obj = ElasticSearchConfig()
        self.request = request

    @property
    def llm(self):
        return ChatOpenAI(model="gpt-4", temperature=0, verbose=True)

    @property
    def base_mapping(self):

        if not self.es_obj.es_url:
            raise es_connect_url_required_exception()
        if not self.es_obj.es_index_name:
            raise es_index_name_required_exception()
        es_client = Elasticsearch(hosts=[self.es_obj.es_url])
        if not es_client.ping():
            raise es_client_err_exception()
        BASE_MAPPING = json.dumps(es_client.indices.get(index=self.es_obj.es_index_name).get(self.es_obj.es_index_name).get("mappings"),
                                  indent=4)
        return BASE_MAPPING

    def get_es_handler(self, question: str) -> str:
        es_promqt_cache_key = 'promqt:redis_es_cache_key_promqt'
        cache_value = self.request.app.state.redis.get(es_promqt_cache_key)
        if cache_value:
            promqt_information = cache_value
        else:
            promqt_information = PROMQT_INFOMATION.format(
                query_dict=question,
                index_name=self.es_obj.es_index_name,
                output_string=OUTPUT_STRING,
                output_string_1=OUTPUT_STRING_1,
                struct_info=self.base_mapping,
            )
            self.request.app.state.redis.set(es_promqt_cache_key, promqt_information)

        result = {}
        while True:
            try:
                response = self.llm.invoke(promqt_information)
                result = json.loads(json.loads(response.json()).get("content", "{}"))
                break
            except json.decoder.JSONDecodeError:
                continue
        return result
